import math

import torch


def cal_mean_covariance(args, model, cal_train_dataset, channel_idx=0):
    """
    计算每个通道预测误差的均值和协方差
    :param args:
    :param model:
    :param cal_train_dataset:
    :param channel_idx:
    :return:
    """
    predictions = []  # 当前窗口内的预测值
    arranged = []  # 整个数据集的多步预测值
    errors = []  # 误差
    with torch.no_grad():
        model.eval()  # eval模式
        last_hidden = model.init_hidden(1)  # 前一个隐藏层
        for t in range(len(cal_train_dataset)):
            # 单步预测
            out, hidden = model.forward(cal_train_dataset[t].unsqueeze(0), last_hidden)  # 前向传播
            predictions.append([])
            arranged.append([])
            errors.append([])
            predictions[t].append(out.data.cpu()[0][0][channel_idx])
            last_hidden = model.repackage_hidden(hidden)  # 更新前一个隐藏层

            # 多步预测
            for _ in range(args.prediction_window_size - 1):
                hidden = model.repackage_hidden(hidden)
                out, hidden = model.forward(out, hidden)
                predictions[t].append(out.data.cpu()[0][0][channel_idx])

            # 使用滑动窗口的思想处理预测值
            if t >= args.prediction_window_size:
                # 将多次多步预测的结果加入
                for step in range(args.prediction_window_size):
                    arranged[t].append(
                        predictions[step + t - args.prediction_window_size][args.prediction_window_size - 1 - step])
                arranged[t] = torch.FloatTensor(arranged[t]).to(args.device)
                errors[t] = arranged[t] - cal_train_dataset[t][0][channel_idx]  # 计算误差
                errors[t] = errors[t].unsqueeze(0)

    errors_tensor = torch.cat(errors[args.prediction_window_size:], dim=0)  # 将误差连接为一个张量
    mean = errors_tensor.mean(dim=0)  # 均值
    covariance = torch.mm(errors_tensor.t(), errors_tensor) / errors_tensor.size(0) - \
                 torch.mm(mean.unsqueeze(1), mean.unsqueeze(0))  # 协方差

    return mean, covariance


def cal_anomaly_score(args, model, dataset, mean, covariance, channel_idx=0):
    """
    计算异常分数值
    :param args:
    :param model:
    :param dataset:
    :param mean:
    :param covariance:
    :param channel_idx:
    :return:
    """
    predictions = []  # 整个窗口内的预测值
    arranged = []  # 整个数据集的多步预测值
    errors = []  # 误差
    with torch.no_grad():
        model.eval()  # eval模式
        last_hidden = model.init_hidden(1)  # 前一个隐藏层
        for t in range(len(dataset)):
            # 单步预测
            out, hidden = model.forward(dataset[t].unsqueeze(0), last_hidden)  # 前向传播
            predictions.append([])
            arranged.append([])
            errors.append([])
            predictions[t].append(out.data.cpu()[0][0][channel_idx])
            last_hidden = model.repackage_hidden(hidden)  # 更新前一个隐藏层

            # 多步预测
            for _ in range(args.prediction_window_size - 1):
                hidden = model.repackage_hidden(hidden)
                out, hidden = model.forward(out, hidden)
                predictions[t].append(out.data.cpu()[0][0][channel_idx])

            # 使用滑动窗口的思想处理预测值
            if t >= args.prediction_window_size:
                # 将多次多步预测的结果加入
                for step in range(args.prediction_window_size):
                    arranged[t].append(
                        predictions[step + t - args.prediction_window_size][args.prediction_window_size - 1 - step])
                arranged[t] = torch.FloatTensor(arranged[t]).to(args.device).unsqueeze(0)
                errors[t] = arranged[t] - dataset[t][0][channel_idx]  # 计算误差
            else:
                arranged[t] = torch.zeros(1, args.prediction_window_size).to(args.device)
                errors[t] = torch.zeros(1, args.prediction_window_size).to(args.device)

    scores = []
    for error in errors:
        # 多元高斯分布模型，计算异常分数值
        cent = error - mean.unsqueeze(0)
        score = torch.mm(torch.mm(cent, torch.inverse(covariance)), cent.t())
        scores.append(score[0][0])

    scores = torch.stack(scores)
    arranged = torch.cat(arranged, dim=0)
    errors = torch.cat(errors, dim=0)

    return scores, arranged, errors


def cal_metrics(args, score, label, steps):
    """
    计算性能指标
    :param args:
    :param score:
    :param label:
    :param steps:
    :return:
    """
    # 计算创建对数均分的一维张量，作为阈值
    threshold = torch.logspace(0, torch.log10(torch.as_tensor(score.max())).item(), steps).to(args.device)

    accuracy = []  # 准确率
    precision = []  # 精确率
    recall = []  # 召回率
    mccs = []  # 马修斯相关系数
    for i in range(len(threshold)):
        anomaly = (score > threshold[i]).float()  # 判断是否为异常
        flag = anomaly * 2 + label
        tn = (flag == 0.0).sum().item()  # tn 检测为非异常，实际也为非异常
        fn = (flag == 1.0).sum().item()  # fn 检测为非异常，实际为异常
        fp = (flag == 2.0).sum().item()  # fp 检测为异常，实际为非异常
        tp = (flag == 3.0).sum().item()  # tp 检测为异常，实际也为异常
        a = (tn + tp) / (tn + tp + fn + fp + 1e-7)  # accuracy
        p = tp / (tp + fp + 1e-7)  # precision
        r = tp / (tp + fn + 1e-7)  # recall
        mcc = (tp * tn - fp * fn) / (math.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)) + 1e-7)

        if a != 0 and p != 0 and r != 0:
            accuracy.append(a)
            precision.append(p)
            recall.append(r)
        mccs.append(mcc)

    # 保存为float型
    accuracy = torch.FloatTensor(accuracy)
    precision = torch.FloatTensor(precision)
    recall = torch.FloatTensor(recall)
    mccs = torch.FloatTensor(mccs)

    # 计算f1-score
    f1_score = 2 * precision * recall / (precision + recall + 1e-7)

    # 确定阈值
    t_i = 0
    for i in range(f1_score.numel()):
        if f1_score[i] == f1_score.max():
            t_i = i
            break
    print('mcc: ', str(mccs[t_i]))

    # 确定异常点
    anomaly_arr = [[]]
    dim = 0
    flag = True
    for i in range(score.numel()):
        if score[i] > threshold[t_i]:
            if flag:
                anomaly_arr[dim].append(i)
                flag = False
                continue
            if i - 1 not in anomaly_arr[dim]:
                anomaly_arr.append([])
                dim += 1
            anomaly_arr[dim].append(i)
    # print(anomaly_arr)

    # anomaly_arr = [[]]
    # dim = 0
    # for i in range(score.numel()):
    #     if score[i] > threshold[t_i]:
    #         anomaly_arr.append([])
    #         dim += 1
    #         anomaly_arr[dim - 1].append(i)

    return accuracy, precision, recall, f1_score, t_i, anomaly_arr
